import codecs
import os
import sys
from collections import OrderedDict

import numpy as np

# reload(sys)
# sys.setdefaultencoding('utf8')

train_dis_prob = open('/dockerdata/rickwwang/project_own/story_generation/data/train.wp_target_500.sentence_pair.discourse').readline().strip()
train_dis_prob = train_dis_prob.split(', ')
train_dis_prob = list(map(lambda x: float(x.split('=')[1]), train_dis_prob))
train_dis_prob = np.asarray(train_dis_prob)


def kl_divergence(p, q):
    return np.sum(np.where(p != 0, p * np.log(p / q), 0))


def format_label(x):
    x = list(map(lambda y: float(y), x))
    max_x = max(x)
    if max_x > 0.8:
        f_x = x.index(max_x)
    else:
        f_x = 8
    return f_x


def extract_discourse(dis_sens, dis_labels):
    dia_data = []
    for ii, (sen_p, sen_d) in enumerate(zip(dis_sens, dis_labels)):
        sid, _, sen1, sen2 = sen_p.strip().split('\t')
        dis_label = sen_d.strip().split('\t')[1:]
        dis_label = format_label(dis_label)
        dia_data.append(dis_label)
    return dia_data


def calc_dis_statstics(dis_data):
    total_sen = len(dis_data)

    metrics = OrderedDict()
    for i in range(9):
        metrics[i] = 0
    for data in dis_data:
        metrics[data] = metrics[data] + 1
    for m in metrics:
        metrics[m] = round(metrics[m] * 1.0 / total_sen * 100, 6)

    return metrics


hyp_path = sys.argv[1]
dis_path = sys.argv[2]
hypos = codecs.open(hyp_path, 'r', encoding='utf8').readlines()
labels = codecs.open(dis_path, 'r', encoding='utf8').readlines()
assert len(hypos) == len(labels)
print('total number of test example', len(hypos))

dis_data = extract_discourse(hypos, labels)
metrics = calc_dis_statstics(dis_data)

predict_prob = np.asarray(list(metrics.values()))
metrics['KL'] = kl_divergence(predict_prob, train_dis_prob)

metrics = list(map(lambda x: str(x[0]) + '=' + str(x[1]), metrics.items()))

print_distinct = ', '.join(metrics)
print('=' * 50)
print(print_distinct)
print('=' * 50)

# write to file
with open(os.path.join(hyp_path + '.discourse'), 'w') as f:
    f.write(print_distinct)
